Covid-19 Analysis

# @title Monthly Covid Incidence by Age
import pandas as pd
import matplotlib.pyplot as plt

# Convert 'date' to datetime if it's not already
covid['date'] = pd.to_datetime(covid['date'])

# Define the age group columns and rename for clarity
age_columns = [
    'confirmed_age_group_0_9_years', 'confirmed_age_group_10_19_years',
    'confirmed_age_group_20_29_years', 'confirmed_age_group_30_39_years',
    'confirmed_age_group_40_49_years', 'confirmed_age_group_50_59_years',
    'confirmed_age_group_60_69_years', 'confirmed_age_group_70_and_older'
]
age_labels = {
    'confirmed_age_group_0_9_years': '0-9',
    'confirmed_age_group_10_19_years': '10-19',
    'confirmed_age_group_20_29_years': '20-29',
    'confirmed_age_group_30_39_years': '30-39',
    'confirmed_age_group_40_49_years': '40-49',
    'confirmed_age_group_50_59_years': '50-59',
    'confirmed_age_group_60_69_years': '60-69',
    'confirmed_age_group_70_and_older': '70+'
}
covid = covid.rename(columns=age_labels)

# Ensure we have monthly data to avoid overcrowding the plot
covid['month'] = covid['date'].dt.to_period('M')
covid_monthly = covid.groupby('month')[list(age_labels.values())].sum()

# Define vaccination start and end dates as periods
vaccination_start = pd.Period('2020-12', freq='M')
vaccination_end = pd.Period('2021-09', freq='M')

# Plotting the stacked bar plot
fig, ax = plt.subplots(figsize=(12, 6))
covid_monthly.plot(kind='bar', stacked=True, ax=ax, colormap="viridis", width=0.8)

# Add labels and formatting
ax.set_title('Monthly COVID-19 Incidence by Age Group')
ax.set_xlabel('Month')
ax.set_ylabel('Number of Confirmed Cases')
ax.legend(title="Age Groups", bbox_to_anchor=(1.05, 1), loc='upper left')

# Calculate x-positions for the vertical lines
start_x = (vaccination_start - covid_monthly.index[0]).n
end_x = (vaccination_end - covid_monthly.index[0]).n

# Add grey vertical lines with adjusted positions
plt.axvline(x=start_x, color='grey', linestyle='--', linewidth=1)
plt.axvline(x=end_x, color='grey', linestyle='--', linewidth=1)

# Add text annotations in the middle of the chart
y_middle = ax.get_ylim()[1] / 2  # Position halfway up the y-axis
ax.text(start_x, y_middle, 'Vaccination Start\n2020-12', color='grey', ha='center', rotation=90)
ax.text(end_x, y_middle, 'Vaccination End\n2021-09', color='grey', ha='center', rotation=90)

# Remove top, right, and left spines
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)

# Only show ticks on bottom and left axis
ax.tick_params(axis='x', bottom=True, top=False)
ax.tick_params(axis='y', left=True, right=False)

# Show plot with adjustments
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# @title COVID-19 Infection Rates by Age Group Over Time
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# Convert 'date' to datetime if it's not already
covid['date'] = pd.to_datetime(covid['date'])

# Define the age group columns and rename for clarity
age_columns = [
    'confirmed_age_group_0_9_years', 'confirmed_age_group_10_19_years',
    'confirmed_age_group_20_29_years', 'confirmed_age_group_30_39_years',
    'confirmed_age_group_40_49_years', 'confirmed_age_group_50_59_years',
    'confirmed_age_group_60_69_years', 'confirmed_age_group_70_and_older'
]
age_labels = {
    'confirmed_age_group_0_9_years': '0-9',
    'confirmed_age_group_10_19_years': '10-19',
    'confirmed_age_group_20_29_years': '20-29',
    'confirmed_age_group_30_39_years': '30-39',
    'confirmed_age_group_40_49_years': '40-49',
    'confirmed_age_group_50_59_years': '50-59',
    'confirmed_age_group_60_69_years': '60-69',
    'confirmed_age_group_70_and_older': '70+'
}
covid = covid.rename(columns=age_labels)

# Ensure we have monthly data to avoid overcrowding the plot
covid['month'] = covid['date'].dt.to_period('M')
covid_monthly = covid.groupby('month')[list(age_labels.values())].sum()

# Define population estimates for 2022 based on provided data
population_estimates = {
    '0-9': 200109,
    '10-19': 192730,
    '20-29': 180063,
    '30-39': 268339,
    '40-49': 250110,
    '50-59': 257933,
    '60-69': 247103,
    '70+': 279370
}

# Calculate infection rates per 100,000 for each age group
infection_rates = covid_monthly.div(population_estimates.values(), axis=1) * 100000

# Define vaccination start and end dates
vaccination_start = pd.Period('2020-12', freq='M')
vaccination_end = pd.Period('2021-09', freq='M')

# Plotting the line plot
fig, ax = plt.subplots(figsize=(12, 6))

# Use the Viridis colormap
colors = plt.cm.viridis(np.linspace(0, 1, len(infection_rates.columns)))

for i, age_group in enumerate(infection_rates.columns):
    ax.plot(infection_rates.index.to_timestamp(), infection_rates[age_group], label=age_group, color=colors[i])

# Add vertical lines for vaccination start and end
start_x = (vaccination_start - covid_monthly.index[0]).n
end_x = (vaccination_end - covid_monthly.index[0]).n
ax.axvline(x=infection_rates.index.to_timestamp()[start_x], color='grey', linestyle='--', linewidth=1)
ax.axvline(x=infection_rates.index.to_timestamp()[end_x], color='grey', linestyle='--', linewidth=1)

# Add text annotations for vaccination start and end
y_middle = ax.get_ylim()[1] / 2  # Position halfway up the y-axis
ax.text(infection_rates.index.to_timestamp()[start_x], y_middle, 'Vaccination Start\n2020-12', color='grey', ha='center', rotation=90)
ax.text(infection_rates.index.to_timestamp()[end_x], y_middle, 'Vaccination End\n2021-09', color='grey', ha='center', rotation=90)

# Add labels and formatting
ax.set_title('COVID-19 Infection Rates by Age Group Over Time')
ax.set_xlabel('Month')
ax.set_ylabel('Infection Rate (per 100,000 population)')
ax.legend(title="Age Groups", bbox_to_anchor=(1.05, 1), loc='upper left')
plt.xticks(rotation=45)

# Remove top, right, and left spines
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)

# Only show ticks on bottom and left axis
ax.tick_params(axis='x', bottom=True, top=False)
ax.tick_params(axis='y', left=True, right=False)

# Show plot with adjustments
plt.tight_layout()
plt.show()

# @title COVID-19 Infection Rates (per 1,000 population) by Age Group and Month
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# Convert 'date' to datetime if it's not already
covid['date'] = pd.to_datetime(covid['date'])

# Define the age group columns and rename for clarity
age_columns = [
    'confirmed_age_group_0_9_years', 'confirmed_age_group_10_19_years',
    'confirmed_age_group_20_29_years', 'confirmed_age_group_30_39_years',
    'confirmed_age_group_40_49_years', 'confirmed_age_group_50_59_years',
    'confirmed_age_group_60_69_years', 'confirmed_age_group_70_and_older'
]
age_labels = {
    'confirmed_age_group_0_9_years': '0-9',
    'confirmed_age_group_10_19_years': '10-19',
    'confirmed_age_group_20_29_years': '20-29',
    'confirmed_age_group_30_39_years': '30-39',
    'confirmed_age_group_40_49_years': '40-49',
    'confirmed_age_group_50_59_years': '50-59',
    'confirmed_age_group_60_69_years': '60-69',
    'confirmed_age_group_70_and_older': '70+'
}
covid = covid.rename(columns=age_labels)

# Ensure we have monthly data to avoid overcrowding the plot
covid['month'] = covid['date'].dt.to_period('M')
covid_monthly = covid.groupby('month')[list(age_labels.values())].sum()

# Define population estimates for 2022 based on provided data
population_estimates = {
    '0-9': 200109,
    '10-19': 192730,
    '20-29': 180063,
    '30-39': 268339,
    '40-49': 250110,
    '50-59': 257933,
    '60-69': 247103,
    '70+': 279370
}

# Calculate infection rates per 1,000 for each age group
infection_rates = covid_monthly.div(population_estimates.values(), axis=1) * 1000

# Reset index for heatmap creation
infection_rates = infection_rates.reset_index()

# Melt the DataFrame for seaborn
heatmap_data = infection_rates.melt(id_vars='month', var_name='Age Group', value_name='Infection Rate')

# Create a pivot table for the heatmap
heatmap_pivot = heatmap_data.pivot(index='Age Group', columns='month', values='Infection Rate')

# Fill NaN values with 0
heatmap_pivot = heatmap_pivot.fillna(0)

# Normalize the heatmap data for column-wise coloring
normed_heatmap = (heatmap_pivot - heatmap_pivot.min()) / (heatmap_pivot.max() - heatmap_pivot.min())

# Filter to start from October 2020 and ensure consistent shapes
normed_heatmap = normed_heatmap.loc[:, '2020-10':'2024-09']  # Adjust according to your data's last month
heatmap_pivot_filtered = heatmap_pivot.loc[:, '2020-10':'2024-09']  # Filter the same way for annotations

# Create a heatmap with proper formatting
plt.figure(figsize=(12, 6))  # Increase figure size for better visibility

# Create the heatmap using normalized infection rates
sns.heatmap(normed_heatmap, cmap='viridis', cbar_kws={'label': 'Normalized Infection Rate', 'shrink': 0.5},  # Adjust shrink for colorbar
            annot=heatmap_pivot_filtered.astype(int).replace(0, '0'), fmt='',  # Use original values for annotations
            linewidths=0.5, linecolor='white', square=True)

# Format x-axis labels to show just year-month
plt.xticks(ticks=range(len(normed_heatmap.columns)), labels=[str(month)[:7] for month in normed_heatmap.columns], rotation=45)

# Set title and labels
plt.title('COVID-19 Infection Rates (per 1,000 population) by Age Group and Month')
plt.xlabel('Month')
plt.ylabel('Age Group')

# Show the plot
plt.tight_layout()
plt.show()

# @title Number of Vaccinated Persons by Vaccination Priority Group and Month
import matplotlib.pyplot as plt
import seaborn as sns

# Dictionary for translating Latvian indications to English
translation_dict = {
    "Persona ar hroniskām slimībām": "Person with chronic illnesses",
    "Persona vecumā virs 60 gadiem": "Person aged over 60 years",
    "Ārstniecības persona": "Medical professional",
    "Cita paaugstinātā riska grupa": "Other high-risk group",
    "Izglītības iestāžu darbinieks": "Educational institution worker",
    "Cits iedzīvotājs": "Other resident",
    "Ārstniecības iestādes darbinieks": "Healthcare facility worker",
    "Kontakts ar personām ar hroniskām slimībām": "Contact with chronic patients",
    "Operatīvo dienestu darbinieks": "Emergency services worker",
    "Citas veselības indikācijas": "Other health indications",
    "Pēc paša vēlēšanās": "By personal choice",
    "Nozaru prioritāro iestāžu darbinieks": "Sector priority institution worker",
    "Grūtniece": "Pregnant person",
    "Vakcinēts ārzemēs": "Vaccinated abroad",
    "Plānveida vakcinācija (pēc vakcinācijas kalendāra)": "Scheduled vaccination (per vaccination calendar)",
    "Imūnsupresīva persona": "Immunosuppressed person",
    "SAC darbinieks": "SAC worker",
    "Ieslodzījumu vietu pārvaldes personāls": "Prison staff",
    "SAC klients": "SAC client",
    "Ieslodzītais": "Prisoner",
    "Epidemioloģiskās indikācijas - kontakts ar infekcijas slimnieku": "Epidemiological indications - contact with infectious person",
    "Speciālo iestāžu klients": "Special institution client",
    "Epidemioloģiskās indikācijas - uzliesmojums, epidēmija": "Epidemiological indications - outbreak, epidemic",
    "Pirms ceļojuma": "Before travel",
    "Arodinfekciju profilakse": "Occupational infection prevention",
    "Veselības indikācijas - trauma, zarnu trakta operācija": "Health indications - trauma, gastrointestinal surgery"
}

# number of vaccinated persons by vaccination indications
grouped = vaccination.groupby(['vaccination_indications', 'month'])['number_of_vaccinated_persons'].sum().reset_index().sort_values(by=['month', 'number_of_vaccinated_persons'], ascending=[True, False])

# Replace Latvian indications with English translations in the DataFrame
grouped['vaccination_indications'] = grouped['vaccination_indications'].replace(translation_dict)

# Calculate the total vaccinations per indication and sort in descending order
indication_order = grouped.groupby("vaccination_indications")['number_of_vaccinated_persons'].sum().sort_values(ascending=False).index

# Pivot the data for heatmap format
heatmap_data = grouped.pivot(index="vaccination_indications", columns="month", values="number_of_vaccinated_persons")

# Reorder the rows in the heatmap data to reflect the sorted indication order
heatmap_data = heatmap_data.reindex(indication_order)

# Plotting the heatmap
plt.figure(figsize=(12, 10))
sns.heatmap(heatmap_data, annot=True, fmt=".0f", cmap="YlGnBu", cbar_kws={'label': 'Number of Vaccinated Persons'})
plt.title("Number of Vaccinated Persons by Vaccination Indications and Month")
plt.xlabel("Month")
plt.ylabel("Vaccination Indications")
plt.xticks(rotation=45)
plt.show()

# @title Monthly Vaccinations by Priority Group for Top Vaccination Indications by Age Group

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

# Assuming we have the full initial DataFrame, which we'll call 'vaccination'
# Step 1: Aggregate data to include 'vaccination_indications', 'month', 'number_of_vaccinated_persons', and 'age_category'
grouped = vaccination.groupby(['vaccination_indications', 'month', 'age_category'], observed=True)['number_of_vaccinated_persons'].sum().reset_index()

# Translate vaccination indications to English
grouped['vaccination_indications'] = grouped['vaccination_indications'].replace(translation_dict)

# Step 2: Determine top 5 indications based on total vaccinations
top_indications = grouped.groupby('vaccination_indications')['number_of_vaccinated_persons'].sum().nlargest(5).index

# Step 3: Filter the data for the top 5 indications and make a copy
filtered_data = grouped[grouped['vaccination_indications'].isin(top_indications)].copy()

# Step 4: Sort age groups in ascending order and create subplots
age_groups = sorted(filtered_data['age_category'].unique())  # Sort age groups in ascending order
num_age_groups = len(age_groups)

# Adjusting plot size for readability
plt.figure(figsize=(12, 16))
palette = sns.color_palette("viridis", len(top_indications))

# Plot each age group in a subplot
for i, age_group in enumerate(age_groups):
    plt.subplot((num_age_groups + 1) // 2, 2, i + 1)  # Arrange subplots in a grid
    age_data = filtered_data[filtered_data['age_category'] == age_group]

    # Pivot data to make 'vaccination_indications' as columns for stacked bars
    pivot_data = age_data.pivot(index='month', columns='vaccination_indications', values='number_of_vaccinated_persons').fillna(0)

    # Filter columns to only include existing top indications for this age group
    relevant_indications = [ind for ind in top_indications if ind in pivot_data.columns]
    pivot_data = pivot_data[relevant_indications]

    # Plot stacked bar chart
    pivot_data.plot(kind='bar', stacked=True, ax=plt.gca(), color=palette[:len(relevant_indications)])

    plt.title(f"Monthly Vaccinations by Priority Group for Age Group: {age_group}")
    plt.xlabel("Month")
    plt.ylabel("No. of Vaccinated Persons")  # Abbreviated y-axis label for clarity
    plt.xticks(rotation=45)
    plt.legend(title="Priority Group", loc='upper left', fontsize='small')

plt.tight_layout()
plt.suptitle("Monthly Vaccinations for Top Vaccination Priority Groups by Age Group", y=1.02)
plt.show()

# @title Cumulative Vaccination Counts Over Time by Age Group for Each Vaccination Stage

import matplotlib.pyplot as plt
import pandas as pd

# Define adjusted population estimates with combined "0-19" group (for reference, but not used in calculations here)
population_estimates = {
    '0-19': 200109 + 192730,  # Combined population for 0-19 age group
    '20-29': 180063,
    '30-39': 268339,
    '40-49': 250110,
    '50-59': 257933,
    '60-69': 247103,
    '70+': 279370
}

# Convert the 'month' column to datetime to ensure compatibility with date filtering
vaccination['month'] = pd.to_datetime(vaccination['month'], errors='coerce')

# Define translations for stages
stage_translations = {
    '1.pote': 'First Dose',
    '2.pote': 'Second Dose',
    '3.pote': 'Third Dose'
}

# Filter stages to include only translated doses, excluding "1.balstvakcinācija"
stages = [stage for stage in vaccination['vaccination_stage'].unique() if stage in stage_translations]
num_stages = len(stages)

# Define a common x-axis range for all plots
common_start_date = pd.to_datetime("2020-12")
common_end_date = pd.to_datetime("2021-09")

# Set up the plot with subplots for each stage, 1 plot per row
fig, axes = plt.subplots(num_stages, 1, figsize=(14, num_stages * 4), sharex=False)

for i, stage in enumerate(stages):
    # Filter data for the current stage and date range
    stage_data = vaccination[(vaccination['vaccination_stage'] == stage) &
                             (vaccination['month'] >= common_start_date) &
                             (vaccination['month'] <= common_end_date)]

    if stage_data.empty:
        print(f"No data available for stage {stage} in the specified date range.")
        continue

    # Aggregate data to get monthly vaccinations by age group
    stage_summary = stage_data.groupby(['age_category', 'month'], observed=True)['number_of_vaccinated_persons'].sum().reset_index()

    # Calculate cumulative vaccinations over time for this stage
    stage_summary['cumulative_vaccinations'] = stage_summary.groupby('age_category', observed=True)['number_of_vaccinated_persons'].cumsum()

    # Pivot data for stacked area chart
    pivot_data = stage_summary.pivot(index='month', columns='age_category', values='cumulative_vaccinations').fillna(0)

    # Plot stacked area chart in the subplot
    pivot_data.plot(kind='area', stacked=True, colormap='viridis', alpha=0.7, ax=axes[i])

    # Set titles and labels for each subplot using English translations
    axes[i].set_title(f"Cumulative Vaccination Counts for {stage_translations[stage]} Over Time by Age Group")
    axes[i].set_ylabel("Cumulative Vaccinations")
    axes[i].set_xlabel("Month")
    axes[i].set_xlim([common_start_date, common_end_date])
    axes[i].legend(title="Age Group", bbox_to_anchor=(1.05, 1), loc='upper left')
    axes[i].tick_params(axis='x', rotation=0)  # Set rotation to 0 for better readability

# Overall title
plt.suptitle("Cumulative Vaccination Counts Over Time by Age Group for Each Vaccination Stage", y=1.02)
plt.tight_layout(rect=[0, 0.04, 1, 0.96])
plt.show()

# @title Monthly Vaccination Counts by Gender for 1st and 2nd Dose (with Cumulative Coverage in %)

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

# Population estimates for cumulative coverage calculation
population_estimates = {
    0: 869057,  # Male population
    1: 1006700  # Female population
}

# Convert 'vaccination_date' to datetime and extract month in 'YYYY-MM' format
vaccination['month'] = pd.to_datetime(vaccination['vaccination_date'], errors='coerce').dt.to_period('M')

# Filter for 1st and 2nd dose data separately
first_dose_data = vaccination[vaccination['vaccination_stage'] == '1.pote']
second_dose_data = vaccination[vaccination['vaccination_stage'] == '2.pote']

# Function to prepare data for plotting
def prepare_data(dose_data, population_estimates):
    # Group by month and gender, summing vaccination counts
    summary = dose_data.groupby(['month', 'vaccinated_person_gender'])['number_of_vaccinated_persons'].sum().reset_index()

    # Calculate cumulative count and coverage
    summary['cumulative_vaccinated'] = summary.groupby('vaccinated_person_gender')['number_of_vaccinated_persons'].cumsum()
    summary['cumulative_coverage'] = summary.apply(
        lambda row: (row['cumulative_vaccinated'] / population_estimates[row['vaccinated_person_gender']]) * 100,
        axis=1
    )

    # Pivot data for heatmap format and cumulative coverage for annotations
    heatmap_data = summary.pivot(index='vaccinated_person_gender', columns='month', values='number_of_vaccinated_persons').fillna(0)
    cumulative_coverage_data = summary.pivot(index='vaccinated_person_gender', columns='month', values='cumulative_coverage').fillna(0)

    # Create annotations with cumulative coverage in parentheses
    annotations = heatmap_data.astype(int).astype(str) + "\n(" + cumulative_coverage_data.round(2).astype(str) + "%)"

    return heatmap_data, annotations

# Prepare data for 1st dose
first_dose_heatmap_data, first_dose_annotations = prepare_data(first_dose_data, population_estimates)

# Prepare data for 2nd dose
second_dose_heatmap_data, second_dose_annotations = prepare_data(second_dose_data, population_estimates)

# Plotting both heatmaps
fig, axes = plt.subplots(2, 1, figsize=(12, 8))

# 1st Dose Heatmap
sns.heatmap(first_dose_heatmap_data, cmap="viridis", annot=first_dose_annotations, fmt="", ax=axes[0], cbar_kws={'label': 'Number of Vaccinated Persons'})
axes[0].set_title("Monthly 1st Dose Vaccination Counts by Gender (with Cumulative Coverage in %)")
axes[0].set_xlabel("Month")
axes[0].set_ylabel("Gender")
axes[0].set_xticklabels(first_dose_heatmap_data.columns.astype(str), rotation=45)
axes[0].set_yticklabels(["Male", "Female"], rotation=0)

# 2nd Dose Heatmap
sns.heatmap(second_dose_heatmap_data, cmap="viridis", annot=second_dose_annotations, fmt="", ax=axes[1], cbar_kws={'label': 'Number of Vaccinated Persons'})
axes[1].set_title("Monthly 2nd Dose Vaccination Counts by Gender (with Cumulative Coverage in %)")
axes[1].set_xlabel("Month")
axes[1].set_ylabel("Gender")
axes[1].set_xticklabels(second_dose_heatmap_data.columns.astype(str), rotation=45)
axes[1].set_yticklabels(["Male", "Female"], rotation=0)

plt.tight_layout()
plt.show()

# @title Monthly Vaccination Counts by Age Group for 1st Dose (Male and Female, with Cumulative Coverage in %)

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

# Aggregated population estimates to match the age categories in the vaccination data
population_estimates_age_gender = {
    '0-19': {'male': 48494 + 55030 + 51205 + 47664, 'female': 45296 + 51289 + 48678 + 45183},
    '20-29': {'male': 42739 + 50490, 'female': 40144 + 46690},
    '30-39': {'male': 68879 + 69622, 'female': 64057 + 65781},
    '40-49': {'male': 61158 + 62894, 'female': 60805 + 65253},
    '50-59': {'male': 61275 + 59520, 'female': 67428 + 69710},
    '60-69': {'male': 59254 + 45217, 'female': 75596 + 67036},
    '70+': {'male': 34981 + 21970, 'female': 61032 + 48034}
}

# Convert 'vaccination_date' to datetime and extract month in 'YYYY-MM' format
vaccination['month'] = pd.to_datetime(vaccination['vaccination_date'], errors='coerce').dt.to_period('M')

# Filter for 1st dose only
first_dose_data = vaccination[vaccination['vaccination_stage'] == '1.pote']

# Function to prepare data for plotting by age group and gender
def prepare_data_age_group(dose_data, gender_label):
    # Filter by gender
    data = dose_data[dose_data['vaccinated_person_gender'] == (0 if gender_label == 'male' else 1)]

    # Group by month and age category, summing vaccination counts
    summary = data.groupby(['month', 'age_category'], observed=True)['number_of_vaccinated_persons'].sum().reset_index()

    # Calculate cumulative count and coverage
    summary['cumulative_vaccinated'] = summary.groupby('age_category', observed=True)['number_of_vaccinated_persons'].cumsum()
    summary['cumulative_coverage'] = summary.apply(
        lambda row: (row['cumulative_vaccinated'] / population_estimates_age_gender[row['age_category']][gender_label]) * 100,
        axis=1
    )

    # Pivot data for heatmap format and cumulative coverage for annotations
    heatmap_data = summary.pivot(index='age_category', columns='month', values='number_of_vaccinated_persons').fillna(0)
    cumulative_coverage_data = summary.pivot(index='age_category', columns='month', values='cumulative_coverage').fillna(0)

    # Create annotations with cumulative coverage in parentheses
    annotations = heatmap_data.astype(int).astype(str) + "\n(" + cumulative_coverage_data.round(2).astype(str) + "%)"

    return heatmap_data, annotations

# Prepare data for 1st dose, males only
male_first_dose_heatmap_data, male_first_dose_annotations = prepare_data_age_group(first_dose_data, 'male')

# Prepare data for 1st dose, females only
female_first_dose_heatmap_data, female_first_dose_annotations = prepare_data_age_group(first_dose_data, 'female')

# Plotting both heatmaps in separate rows
fig, axes = plt.subplots(2, 1, figsize=(12, 10))

# Male Heatmap
sns.heatmap(male_first_dose_heatmap_data, cmap="viridis", annot=male_first_dose_annotations, fmt="", ax=axes[0], cbar_kws={'label': 'Number of Vaccinated Persons'})
axes[0].set_title("Monthly 1st Dose Vaccination Counts by Age Group (Males, with Cumulative Coverage in %)")
axes[0].set_xlabel("Month")
axes[0].set_ylabel("Age Group")
axes[0].set_xticklabels(male_first_dose_heatmap_data.columns.astype(str), rotation=45)

# Female Heatmap
sns.heatmap(female_first_dose_heatmap_data, cmap="viridis", annot=female_first_dose_annotations, fmt="", ax=axes[1], cbar_kws={'label': 'Number of Vaccinated Persons'})
axes[1].set_title("Monthly 1st Dose Vaccination Counts by Age Group (Females, with Cumulative Coverage in %)")
axes[1].set_xlabel("Month")
axes[1].set_ylabel("Age Group")
axes[1].set_xticklabels(female_first_dose_heatmap_data.columns.astype(str), rotation=45)

plt.tight_layout()
plt.show()

# @title Monthly Number of Vaccinated Persons by Preparation Type

import matplotlib.pyplot as plt
from matplotlib.cm import viridis
import pandas as pd

# Mapping dictionary
preparation_mapping = {
    "BBIBP-CorV (Sinopharm)": "Sinopharm",
    "Comirnaty": "Pfizer-BioNTech",
    "CoronaVac(Sinovac)": "Sinovac",
    "Covshield(ChAdOx1_nCov-19)": "AstraZeneca",
    "Jcovden": "J&J Janssen",
    "Spikevax": "Moderna",
    "Vaxzevria": "AstraZeneca"
}
# Apply the mapping to replace technical names with commonly known names
vaccination['preparation'] = vaccination['preparation'].replace(preparation_mapping)

# Create a 'year_month' column from 'vaccination_date' in 'YYYY-MM' format
vaccination['year_month'] = pd.to_datetime(vaccination['vaccination_date'], errors='coerce').dt.to_period('M')

# Group by year_month and the updated preparation column
monthly_vaccination_counts = vaccination.groupby(['year_month', 'preparation'])['number_of_vaccinated_persons'].sum().unstack(fill_value=0)

# Generate colors from the Viridis palette
num_preparations = len(monthly_vaccination_counts.columns)
colors = [viridis(i / num_preparations) for i in range(num_preparations)]

# Plotting with adjusted figure size for a smaller height
fig, ax = plt.subplots(figsize=(12, 6))
monthly_vaccination_counts.plot(kind='bar', stacked=True, color=colors, width=0.8, ax=ax)

# Apply chart formatting
ax.set_title('Monthly Number of Vaccinated Persons by Brand Vaccine Names')
ax.set_xlabel('Year-Month')
ax.set_ylabel('Number of Vaccinated Persons')

# Display ticks for each month
ax.set_xticks(range(len(monthly_vaccination_counts.index)))
ax.set_xticklabels([str(period) for period in monthly_vaccination_counts.index], rotation=45)

# Adjust legend
plt.legend(title='Known Vaccine Names', bbox_to_anchor=(1.05, 1), loc='upper left', fontsize='small', title_fontsize='medium')

# Customize the chart by hiding all spines except the bottom
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(True)

plt.tight_layout()
plt.show()

# @title Proportional Stacked Bar Plot of Vaccine Choice Within Each Age Group

import matplotlib.pyplot as plt
import pandas as pd

# Apply the mapping to replace technical names with commonly known names
vaccination['preparation'] = vaccination['preparation'].replace(preparation_mapping)

# Group by age_category and preparation to get the total number of vaccinated persons
age_vaccine_counts = vaccination.groupby(['age_category', 'preparation'], observed=False)['number_of_vaccinated_persons'].sum().unstack(fill_value=0)

# Normalize the counts within each age group to get proportions
age_vaccine_proportions = age_vaccine_counts.div(age_vaccine_counts.sum(axis=1), axis=0)

# Generate colors from the Viridis palette
num_preparations = len(age_vaccine_proportions.columns)
colors = [viridis(i / num_preparations) for i in range(num_preparations)]

# Plotting
fig, ax = plt.subplots(figsize=(12, 8))
age_vaccine_proportions.plot(kind='bar', stacked=True, color=colors, ax=ax)

# Apply chart formatting
ax.set_title('Proportion of Vaccine Preparations by Age Category')
ax.set_xlabel('Age Category')
ax.set_ylabel('Proportion of Vaccinated Persons')

# Format y-axis as percentages with a tick for 100%
ax.yaxis.set_major_formatter(plt.matplotlib.ticker.PercentFormatter(1.0))
ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])  # Adds 100% (1.0) as a y-tick

# Adjust legend
plt.legend(title='Vaccine Type', bbox_to_anchor=(1.05, 1), loc='upper left', fontsize='small', title_fontsize='medium')

# Customize the chart by hiding all spines except the bottom
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(True)

plt.tight_layout()
plt.show()

# @title Monthly Vaccination Breakdown by Preparation Type for Each Age Category

import matplotlib.pyplot as plt
import pandas as pd

# Apply the mapping to replace technical names with commonly known names
vaccination['preparation'] = vaccination['preparation'].replace(preparation_mapping)

# Ensure 'vaccination_date' is in datetime format and extract month
vaccination['vaccination_date'] = pd.to_datetime(vaccination['vaccination_date'])
vaccination['month'] = vaccination['vaccination_date'].dt.to_period('M')

# Group by month, age_category, and preparation to get the total number of vaccinated persons, with observed=False to suppress warning
monthly_age_vaccine_counts = vaccination.groupby(['month', 'age_category', 'preparation'], observed=False)['number_of_vaccinated_persons'].sum().unstack(fill_value=0)

# Set up subplots for each age category
age_categories = monthly_age_vaccine_counts.index.get_level_values('age_category').unique()
num_age_categories = len(age_categories)

fig, axes = plt.subplots(num_age_categories, 1, figsize=(14, 4 * num_age_categories), sharex=True)

# Generate colors from the Viridis palette
num_preparations = len(monthly_age_vaccine_counts.columns)
colors = [viridis(i / num_preparations) for i in range(num_preparations)]

# Plotting for each age category
for i, age_category in enumerate(age_categories):
    ax = axes[i]
    data = monthly_age_vaccine_counts.xs(age_category, level='age_category')
    data.plot(kind='bar', stacked=True, color=colors, ax=ax, width=0.8)

    # Formatting each subplot
    ax.set_title(f'Age Category: {age_category}')
    ax.set_ylabel('Number of Vaccinated Persons')
    ax.legend(title='Vaccine Type', fontsize='small', title_fontsize='medium', loc='upper left')

    # Adjust y-axis to automatically scale based on the data
    ax.get_yaxis().set_major_formatter(plt.matplotlib.ticker.ScalarFormatter())
    # Customize the chart by hiding all spines except the bottom
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['bottom'].set_visible(True)

# Set common labels and tighten layout
axes[-1].set_xlabel('Month')
plt.tight_layout()
plt.show()

# @title COVID-19 Vaccination Coverage for First Dose Visualization

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib.dates as mdates

# Assuming 'vaccination' dataframe is provided

# Step 1: Filter for '1.pote' in the vaccination stage and remove NaNs in age_category
vaccination_first_dose = vaccination[vaccination['vaccination_stage'] == '1.pote'].copy()
vaccination_first_dose = vaccination_first_dose.dropna(subset=['age_category'])
vaccination_first_dose['age_category'] = vaccination_first_dose['age_category'].astype(str)

# Filter out any remaining unexpected NaN values in age_category directly
vaccination_first_dose = vaccination_first_dose[vaccination_first_dose['age_category'] != 'nan']

# Aggregate the number of vaccinated persons by age_category and month
vaccination_agg = vaccination_first_dose.groupby(['month', 'age_category']).agg(
    num_vaccinated=('number_of_vaccinated_persons', 'sum')
).reset_index()

# Define population estimates and map to age categories
population_estimates_age_gender = {
    '0-19': {'male': 48494 + 55030 + 51205 + 47664, 'female': 45296 + 51289 + 48678 + 45183},
    '20-29': {'male': 42739 + 50490, 'female': 40144 + 46690},
    '30-39': {'male': 68879 + 69622, 'female': 64057 + 65781},
    '40-49': {'male': 61158 + 62894, 'female': 60805 + 65253},
    '50-59': {'male': 61275 + 59520, 'female': 67428 + 69710},
    '60-69': {'male': 59254 + 45217, 'female': 75596 + 67036},
    '70+': {'male': 34981 + 21970, 'female': 61032 + 48034}
}
population_totals = {age: sum(gender.values()) for age, gender in population_estimates_age_gender.items()}
vaccination_agg['estimated_population'] = vaccination_agg['age_category'].map(population_totals)

# Convert 'month' to datetime format and calculate cumulative coverage
vaccination_agg['month'] = pd.to_datetime(vaccination_agg['month'].astype(str))
vaccination_agg = vaccination_agg.sort_values(by=['age_category', 'month'])
vaccination_agg['cumulative_vaccinated'] = vaccination_agg.groupby('age_category')['num_vaccinated'].cumsum()
vaccination_agg['coverage_per_cent'] = (vaccination_agg['cumulative_vaccinated'] / vaccination_agg['estimated_population']) * 100

# Prepare the vaccination coverage result table
result_table = vaccination_agg[['month', 'age_category', 'coverage_per_cent']]

# Step 4: Set up plot style
sns.set_palette(sns.color_palette("viridis", n_colors=result_table['age_category'].nunique()))

# Initialize the plot
fig, ax1 = plt.subplots(figsize=(12, 6))

# Plot vaccination coverage for each age category on the primary y-axis
sns.lineplot(data=result_table, x='month', y='coverage_per_cent', hue='age_category', marker='o', ax=ax1)
ax1.set_ylim(0, 100)  # Set y-axis limits for coverage

# Adding herd immunity threshold lines and slightly shifting text to the right
herd_immunity_thresholds = [60, 70, 85]
for threshold in herd_immunity_thresholds:
    ax1.axhline(y=threshold, color='gray', linestyle='--', linewidth=1)
    ax1.text(result_table['month'].min() + pd.Timedelta(days=10), threshold + 1, f'{threshold}%', color='gray')

# Customize the plot
ax1.set_title('COVID-19 Vaccination Coverage for First Dose by Age Category')
ax1.set_xlabel('Month')
ax1.set_ylabel('Vaccination Coverage (%)')
ax1.legend(title='Age Category', loc='upper left', bbox_to_anchor=(1, 1))  # Move legend outside plot area
ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))

# Remove grid, ensure no gridlines appear on both axes
ax1.grid(False)

# Hide all but bottom spines as specified
for spine in ['top', 'right', 'left']:
    ax1.spines[spine].set_visible(False)

plt.tight_layout()
plt.show()

# @title COVID-19 Vaccination Coverage for Second Dose Visualization

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib.dates as mdates

# Assuming 'vaccination' dataframe is provided

# Step 1: Filter for '2.pote' in the vaccination stage and remove NaNs in age_category
vaccination_second_dose = vaccination[vaccination['vaccination_stage'] == '2.pote'].copy()
vaccination_second_dose = vaccination_second_dose.dropna(subset=['age_category'])
vaccination_second_dose['age_category'] = vaccination_second_dose['age_category'].astype(str)

# Filter out any remaining unexpected NaN values in age_category directly
vaccination_second_dose = vaccination_second_dose[vaccination_second_dose['age_category'] != 'nan']

# Aggregate the number of vaccinated persons by age_category and month
vaccination_agg = vaccination_second_dose.groupby(['month', 'age_category']).agg(
    num_vaccinated=('number_of_vaccinated_persons', 'sum')
).reset_index()

# Define population estimates and map to age categories
population_estimates_age_gender = {
    '0-19': {'male': 48494 + 55030 + 51205 + 47664, 'female': 45296 + 51289 + 48678 + 45183},
    '20-29': {'male': 42739 + 50490, 'female': 40144 + 46690},
    '30-39': {'male': 68879 + 69622, 'female': 64057 + 65781},
    '40-49': {'male': 61158 + 62894, 'female': 60805 + 65253},
    '50-59': {'male': 61275 + 59520, 'female': 67428 + 69710},
    '60-69': {'male': 59254 + 45217, 'female': 75596 + 67036},
    '70+': {'male': 34981 + 21970, 'female': 61032 + 48034}
}
population_totals = {age: sum(gender.values()) for age, gender in population_estimates_age_gender.items()}
vaccination_agg['estimated_population'] = vaccination_agg['age_category'].map(population_totals)

# Convert 'month' to datetime format and calculate cumulative coverage
vaccination_agg['month'] = pd.to_datetime(vaccination_agg['month'].astype(str))
vaccination_agg = vaccination_agg.sort_values(by=['age_category', 'month'])
vaccination_agg['cumulative_vaccinated'] = vaccination_agg.groupby('age_category')['num_vaccinated'].cumsum()
vaccination_agg['coverage_per_cent'] = (vaccination_agg['cumulative_vaccinated'] / vaccination_agg['estimated_population']) * 100

# Prepare the vaccination coverage result table
result_table = vaccination_agg[['month', 'age_category', 'coverage_per_cent']]

# Step 4: Set up plot style
sns.set_palette(sns.color_palette("viridis", n_colors=result_table['age_category'].nunique()))

# Initialize the plot
fig, ax1 = plt.subplots(figsize=(12, 6))

# Plot vaccination coverage for each age category on the primary y-axis
sns.lineplot(data=result_table, x='month', y='coverage_per_cent', hue='age_category', marker='o', ax=ax1)
ax1.set_ylim(0, 100)  # Set y-axis limits for coverage

# Adding herd immunity threshold lines and slightly shifting text to the right
herd_immunity_thresholds = [60, 70, 85]
for threshold in herd_immunity_thresholds:
    ax1.axhline(y=threshold, color='gray', linestyle='--', linewidth=1)
    ax1.text(result_table['month'].min() + pd.Timedelta(days=10), threshold + 1, f'{threshold}%', color='gray')

# Customize the plot
ax1.set_title('COVID-19 Vaccination Coverage for Second Dose')
ax1.set_xlabel('Month')
ax1.set_ylabel('Vaccination Coverage (%)')
ax1.legend(title='Age Category', loc='upper left', bbox_to_anchor=(1, 1))  # Move legend outside plot area
ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))

# Remove grid, ensure no gridlines appear on both axes
ax1.grid(False)

# Hide all but bottom spines as specified
for spine in ['top', 'right', 'left']:
    ax1.spines[spine].set_visible(False)

plt.tight_layout()
plt.show()

# @title COVID-19 Vaccination Coverage for Second Dose with Recovery Index
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib.dates as mdates

# Step 1: Prepare output_table for recovery index
# Assuming 'covid' DataFrame is available

# Ensure 'date' is in datetime format
covid['date'] = pd.to_datetime(covid['date'], errors='coerce')

# Extract year-month for grouping
covid['year_month'] = covid['date'].dt.to_period('M')

# Aggregate to get the last value of cumulative recovered cases each month
# and calculate the cumulative sum for confirmed cases
monthly_summary = covid.groupby('year_month').agg(
    cumulative_recovered_cases=('cumulative_recovered_cases', 'last'),
    total_cases=('confirmed_covid19_cases', 'sum')
).reset_index()

# Convert 'year_month' to datetime format if it is in Period format
if monthly_summary['year_month'].dtype == 'period[M]':
    monthly_summary['month'] = monthly_summary['year_month'].dt.to_timestamp()

# Calculate cumulative sum for total cases
monthly_summary['cumulative_total_cases'] = monthly_summary['total_cases'].cumsum()

# Calculate recovery index as a percentage
monthly_summary['recovery_index'] = (monthly_summary['cumulative_recovered_cases'] / monthly_summary['cumulative_total_cases']) * 100

# Define the vaccination period
vaccination_period_start = pd.to_datetime("2020-12-01")
vaccination_period_end = pd.to_datetime("2021-09-30")

# Filter the table for the vaccination period
output_table = monthly_summary[(monthly_summary['month'] >= vaccination_period_start) &
                               (monthly_summary['month'] <= vaccination_period_end)].reset_index(drop=True)

# Step 2: Prepare vaccination coverage data for the same period
# Assuming 'vaccination' DataFrame contains 'vaccination_stage', 'month', 'age_category', and 'number_of_vaccinated_persons'

# Convert 'month' to datetime format if it is an object or Period format
if vaccination['month'].dtype == 'object':
    vaccination['month'] = pd.to_datetime(vaccination['month'], errors='coerce')
elif vaccination['month'].dtype == 'period[M]':
    vaccination['month'] = vaccination['month'].dt.to_timestamp()

# Filter the vaccination data for the second dose within the vaccination period
vaccination_second_dose = vaccination[(vaccination['vaccination_stage'] == '2.pote') &
                                      (vaccination['month'] >= vaccination_period_start) &
                                      (vaccination['month'] <= vaccination_period_end)].copy()

# Remove NaNs in age_category and convert to string
vaccination_second_dose = vaccination_second_dose.dropna(subset=['age_category'])
vaccination_second_dose['age_category'] = vaccination_second_dose['age_category'].astype(str)

# Aggregate the number of vaccinated persons by age_category and month
vaccination_agg = vaccination_second_dose.groupby(['month', 'age_category']).agg(
    num_vaccinated=('number_of_vaccinated_persons', 'sum')
).reset_index()

# Define population estimates and map to age categories
population_estimates_age_gender = {
    '0-19': {'male': 48494 + 55030 + 51205 + 47664, 'female': 45296 + 51289 + 48678 + 45183},
    '20-29': {'male': 42739 + 50490, 'female': 40144 + 46690},
    '30-39': {'male': 68879 + 69622, 'female': 64057 + 65781},
    '40-49': {'male': 61158 + 62894, 'female': 60805 + 65253},
    '50-59': {'male': 61275 + 59520, 'female': 67428 + 69710},
    '60-69': {'male': 59254 + 45217, 'female': 75596 + 67036},
    '70+': {'male': 34981 + 21970, 'female': 61032 + 48034}
}
population_totals = {age: sum(gender.values()) for age, gender in population_estimates_age_gender.items()}
vaccination_agg['estimated_population'] = vaccination_agg['age_category'].map(population_totals)

# Convert 'month' to datetime format if it is not already, then calculate cumulative coverage
vaccination_agg = vaccination_agg.sort_values(by=['age_category', 'month'])
vaccination_agg['cumulative_vaccinated'] = vaccination_agg.groupby('age_category')['num_vaccinated'].cumsum()
vaccination_agg['coverage_per_cent'] = (vaccination_agg['cumulative_vaccinated'] / vaccination_agg['estimated_population']) * 100

# Prepare the vaccination coverage result table
result_table = vaccination_agg[['month', 'age_category', 'coverage_per_cent']]

# Step 3: Plot vaccination coverage and recovery index

# Initialize the plot
fig, ax1 = plt.subplots(figsize=(12, 6))

# Plot vaccination coverage for each age category on the primary y-axis using viridis color palette
sns.lineplot(data=result_table, x='month', y='coverage_per_cent', hue='age_category', marker='o', palette='viridis', ax=ax1)
ax1.set_ylim(0, 100)  # Set y-axis limits for coverage

# Adding herd immunity threshold lines and slightly shifting text to the right
herd_immunity_thresholds = [60, 70, 85]
for threshold in herd_immunity_thresholds:
    ax1.axhline(y=threshold, color='gray', linestyle='--', linewidth=1)
    ax1.text(result_table['month'].min() + pd.Timedelta(days=10), threshold + 1, f'{threshold}%', color='gray')

# Customize the primary y-axis plot
ax1.set_title('COVID-19 Vaccination Coverage for Second Dose with Recovery Index')
ax1.set_xlabel('Month')
ax1.set_ylabel('Vaccination Coverage (%)')
# Move the legend further outside to avoid overlap with the secondary y-axis ticks
ax1.legend(title='Age Category', loc='upper left', bbox_to_anchor=(1.05, 1))
ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))

# Secondary y-axis for recovery index with a black dashed line
ax2 = ax1.twinx()
sns.lineplot(data=output_table, x='month', y='recovery_index', color='black', linestyle='--', marker='x', ax=ax2, label='Recovery Index (%)')
ax2.set_ylabel('Recovery Index (%)')
ax2.tick_params(axis='y')  # Default color for secondary y-axis ticks

# Remove top, left, and right spines
for spine in ['top', 'right', 'left']:
    ax1.spines[spine].set_visible(False)
    ax2.spines[spine].set_visible(False)

plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# @title Vaccinated Persons by Facility and Age Category
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import warnings

# Silence warnings
warnings.filterwarnings("ignore", category=FutureWarning)

# Assuming your vaccination data is in a DataFrame named 'vaccination'

# Step 1: Select the top 25 facilities by vaccination counts and clean facility names
top_25_facilities = vaccination['vaccination_facility_name'].value_counts().nlargest(25).index
vaccination['facility_group'] = vaccination['vaccination_facility_name'].apply(lambda x: x.split(',')[0] if x in top_25_facilities else None)

# Filter out rows where 'facility_group' is None (i.e., "Other" facilities)
filtered_vaccination = vaccination.dropna(subset=['facility_group'])

# Step 2: Group by the cleaned facility name and age category to get the sum of vaccinated persons
facility_age_data = filtered_vaccination.groupby(['facility_group', 'age_category'], observed=True)['number_of_vaccinated_persons'].sum().unstack(fill_value=0)

# Step 3: Order facilities by total vaccinated persons in descending order
facility_age_data['Total_Vaccinated'] = facility_age_data.sum(axis=1)
facility_age_data = facility_age_data.sort_values('Total_Vaccinated', ascending=False).drop(columns='Total_Vaccinated')

# Step 4: Normalize each row for row-wise color intensity
facility_age_data_normalized = facility_age_data.div(facility_age_data.max(axis=1), axis=0)

# Step 5: Plot the heatmap using viridis colormap, which is row-normalized
plt.figure(figsize=(12, 10))
sns.heatmap(facility_age_data_normalized, cmap="viridis", annot=facility_age_data, fmt="d",
            cbar_kws={'label': 'Normalized Vaccinated Persons'}, linewidths=.5)
plt.title("Number of Vaccinated Persons by Facility and Age Category")
plt.xlabel("Age Category")
plt.ylabel("Vaccination Facility (Top 25, Ordered)")
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# @title Vaccinated Persons by Facility and Month
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import warnings

# Silence warnings
warnings.filterwarnings("ignore", category=FutureWarning)

# Assuming your vaccination data is in a DataFrame named 'vaccination'

# Convert 'month' column to datetime format if it isn't already, then format it to YYYY-MM
vaccination['month'] = pd.to_datetime(vaccination['month']).dt.strftime('%Y-%m')

# Step 1: Select the top 25 facilities by vaccination counts and clean facility names
top_25_facilities = vaccination['vaccination_facility_name'].value_counts().nlargest(25).index
vaccination['facility_group'] = vaccination['vaccination_facility_name'].apply(lambda x: x.split(',')[0] if x in top_25_facilities else None)

# Filter out rows where 'facility_group' is None (i.e., "Other" facilities)
filtered_vaccination = vaccination.dropna(subset=['facility_group'])

# Step 2: Group by the cleaned facility name and formatted month to get the sum of vaccinated persons
facility_month_data = filtered_vaccination.groupby(['facility_group', 'month'], observed=True)['number_of_vaccinated_persons'].sum().unstack(fill_value=0)

# Step 3: Order facilities by total vaccinated persons in descending order
facility_month_data['Total_Vaccinated'] = facility_month_data.sum(axis=1)
facility_month_data = facility_month_data.sort_values('Total_Vaccinated', ascending=False).drop(columns='Total_Vaccinated')

# Step 4: Normalize each row for row-wise color intensity
facility_month_data_normalized = facility_month_data.div(facility_month_data.max(axis=1), axis=0)

# Step 5: Plot the heatmap using viridis colormap, which is row-normalized
plt.figure(figsize=(12, 10))
sns.heatmap(facility_month_data_normalized, cmap="viridis", annot=facility_month_data, fmt="d",
            cbar_kws={'label': 'Normalized Vaccinated Persons'}, linewidths=.5)
plt.title("Number of Vaccinated Persons by Facility and Month")
plt.xlabel("Month")
plt.ylabel("Vaccination Facility (Top 25, Ordered)")
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()